import sys,re,subprocess,time,datetime,yaml,os
from termcolor import colored
import open3d #don't remove. for check env
CONDA_ACTIVATE_CMD = "conda activate 3dpidm"
def check_tmux_time(session_name,command):
    result = subprocess.run(f"tmux ls | grep '{session_name}: 1 windows'", shell=True, capture_output=True, text=True)
    if result.returncode != 0 or not result.stdout.strip():
        return
    created_strs = result.stdout.strip().split('created ')[1].split(')')
    created_str = created_strs[0]
    session_created = datetime.datetime.strptime(created_str, '%a %b %d %H:%M:%S %Y').timestamp()
    current_time = time.time()
    time_diff = current_time - session_created
    if time_diff > 15*60:
        user_input = input(colored(f"Session {session_name} already running for {time.strftime('%H:%M:%S', time.gmtime(time_diff))}","red",attrs=["bold"])+f" created at {created_str},\nPlease first run: "+colored(f"tmux attach -t {session_name}","blue",attrs=["bold"])+f" and confirm execute command:\n{command} (yes/no)")
        assert user_input.lower() == "yes", f"User cancelled operation for session {session_name}"
def log_command_execution(commands_data):
    timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
    log_data = {'argv': ' '.join(sys.argv)}
    log_data.update(commands_data)
    log_file = f"{sys.argv[1]}_log.yaml"
    if os.path.exists(log_file):
        with open(log_file, 'r') as f:
            existing_log = yaml.safe_load(f) or {}
    else:
        existing_log = {}
    existing_log['command_log'] = existing_log.get('command_log', {})
    existing_log['command_log'][timestamp] = log_data
    with open(log_file, 'w') as f:
        yaml.dump(existing_log, f, default_flow_style=False)
def parse_lists_from_argv(argv):
    lists = []
    list_positions = []
    for i, arg in enumerate(argv):
        if arg.startswith('[') and arg.endswith(']'):
            try:
                content = arg[1:-1]
                items = [item.strip() for item in content.split(',')]
                parsed_list = []
                for item in items:
                    try:
                        parsed_list.append(int(item))
                    except ValueError:
                        parsed_list.append(item)
                lists.append(parsed_list)
                list_positions.append(i)
            except:
                assert False, f"Failed to parse list: {arg}"
    return lists, list_positions
def generate_session_name(params,use_perm):
    # if not use_perm:
    #     return str(params[0])
    if len(params) == 1:
        return 'G'+str(params[0])
    return 'G'+'-'.join(map(str, params))
def generate_commands(pipeline_name, argv, lists, list_positions, use_perm=False, original_argv=None):
    if not lists:
        return {"sessions": {}, "cmd_type": "run"}
    if not use_perm:
        assert all(len(lst) == len(lists[0]) for lst in lists), f"All lists must have same length for default no-perm mode, got lengths: {[len(lst) for lst in lists]}"
        combinations = []
        for i in range(len(lists[0])):
            combo = [lst[i] for lst in lists]
            combinations.append(combo)
    else:
        def cartesian_product(arrays):
            if len(arrays) == 1:
                return [[x] for x in arrays[0]]
            result = []
            for item in arrays[0]:
                for sub_product in cartesian_product(arrays[1:]):
                    result.append([item] + sub_product)
            return result
        combinations = cartesian_product(lists)
    sessions = {}
    cmd_type = "run"
    if '--int' in original_argv:
        cmd_type = "int"
    elif '--kill' in original_argv:
        cmd_type = "kill"
    elif '--rerun' in original_argv:
        cmd_type = "rerun"
    for combo in combinations:
        session_name = generate_session_name(combo,use_perm)
        new_argv = argv.copy()
        for i, pos in enumerate(list_positions):
            new_argv[pos] = str(combo[i])
        param_str = ' '.join(new_argv)
        if cmd_type == "int":
            command = f'tmux send-keys -t {session_name} C-c'
            check_tmux_time(session_name,command)
        elif cmd_type == "kill":
            command = f'tmux send-keys -t {session_name} C-c; tmux kill-session -t {session_name}'
            check_tmux_time(session_name,command)
        elif cmd_type == "rerun":
            command = f'tmux send-keys -t {session_name} C-c; tmux kill-session -t {session_name}; tmux new -d -s {session_name} "bash"; tmux send-keys -t {session_name} "{CONDA_ACTIVATE_CMD}; python3 class_run_fullpipeline.py {pipeline_name} {param_str}" Enter; tmux send-keys -t {session_name} "tmux kill-session -t {session_name}"'
            check_tmux_time(session_name,command)
        else:
            command = f'tmux new -d -s {session_name} "bash"; tmux send-keys -t {session_name} "{CONDA_ACTIVATE_CMD}; python3 class_run_fullpipeline.py {pipeline_name} {param_str}" Enter; tmux send-keys -t {session_name} "tmux kill-session -t {session_name}"'
        attach_cmd = f'tmux attach -t {session_name}'
        sessions[session_name] = {"cmd": command, "attach_cmd": attach_cmd}
    return {"sessions": sessions, "cmd_type": cmd_type}
pipeline_name = sys.argv[1]
argv_without_pipeline = sys.argv[2:].copy()
use_perm = '--perm' in argv_without_pipeline

timer_seconds = None
timer_index = -1
for i, arg in enumerate(argv_without_pipeline):
    if arg == '--timer' and i + 1 < len(argv_without_pipeline):
        timer_seconds = int(argv_without_pipeline[i + 1])
        timer_index = i
        break
argv_without_pipeline = [arg for arg in argv_without_pipeline if arg not in ['--perm','--int', '--kill', '--rerun', '--timer']]
timer_index != -1 and timer_index + 1 < len(sys.argv[2:]) and argv_without_pipeline.remove(str(timer_seconds))
lists, list_positions = parse_lists_from_argv(argv_without_pipeline)
if not lists:
    command_args = ["python3", "class_run_fullpipeline.py", pipeline_name] + argv_without_pipeline
    print(f"Running directly: {' '.join(command_args)}")
    subprocess.run(command_args)
    log_command_execution({'direct_run': ' '.join(command_args)})
else:
    commands_data = generate_commands(pipeline_name, argv_without_pipeline, lists, list_positions, use_perm, sys.argv[2:])
    if commands_data["cmd_type"] in ["int", "kill"]:
        input("confirm execute SIGINT?")
    log_command_execution(commands_data)
    for session_name, session_data in commands_data["sessions"].items():
        print(session_data["cmd"])
        subprocess.run(session_data["cmd"], shell=True)
        print(session_data["attach_cmd"])
        if commands_data["cmd_type"] in ["run", "rerun"]:
            if timer_seconds is not None:
                print(f"Waiting {timer_seconds} seconds before continuing...")
                time.sleep(timer_seconds)
            else:
                if input("continue?"):
                    continue
